import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # 自动初始化 CUDA 上下文
import numpy as np
import time

# -------------------------- 1. 初始化 TensorRT 组件 --------------------------
TRT_LOGGER = trt.Logger(trt.Logger.INFO)  # 日志级别（WARNING 减少冗余输出）
builder = trt.Builder(TRT_LOGGER)
print('build end。。。')
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(EXPLICIT_BATCH)  # 显式批处理（支持动态 batch）
config = builder.create_builder_config()

# -------------------------- 2. 构建网络结构（关键：必须添加输入/输出/层） --------------------------
# 示例：构建一个简单的全连接网络（输入: (1,256) → 输出: (1,10)）
# 2.1 添加输入层（名称、数据类型、形状）
input_tensor = network.add_input(
    name="input",
    dtype=trt.float32,
    shape=(1, 256)  # 显式批处理：batch_size=1，特征维度=256
)

# 2.2 添加全连接层（输入张量、输出维度、权重/偏置） trt.Weights.from_numpy(arr)
# 注意：权重/偏置需是 trt.Weights 类型（示例用随机数，实际需替换为模型权重）
fc_weights = trt.Weights.from_numpy(
    trt.float32,
    np.random.randn(10, 256).astype(np.float32)  # 全连接层权重（输出维度×输入维度）
)
fc_biases = trt.Weights(
    trt.float32,
    np.random.randn(10).astype(np.float32)       # 全连接层偏置（输出维度）
)
fc_layer = network.add_fully_connected(
    input=input_tensor,
    num_outputs=10,
    weights=fc_weights,
    biases=fc_biases
)

# 2.3 添加输出层（标记为网络输出）
output_tensor = fc_layer.get_output(0)
output_tensor.name = "output"
network.mark_output(output_tensor)

# -------------------------- 3. 编译并反序列化引擎 --------------------------
# 3.1 编译网络为序列化字节流（需 network + config）
serialized_engine = builder.build_serialized_network(network, config)
if not serialized_engine:
    raise RuntimeError("Failed to build serialized engine!")

# 3.2 反序列化字节流为可执行引擎
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(serialized_engine)
if not engine:
    raise RuntimeError("Failed to deserialize engine!")

# -------------------------- 4. 准备推理上下文和内存 --------------------------
# 4.1 创建推理上下文（执行引擎的实例）
context = engine.create_execution_context()

# 4.2 准备输入/输出数据（主机内存）
batch_size = 1
input_shape = (batch_size, 256)
output_shape = (batch_size, 10)

input_data = np.random.randn(*input_shape).astype(np.float32)  # 随机输入
output_data = np.empty(output_shape, dtype=np.float32)         # 空数组存输出

# 4.3 分配设备内存（CUDA 内存）
d_input = cuda.mem_alloc(input_data.nbytes)   # 输入设备内存
d_output = cuda.mem_alloc(output_data.nbytes) # 输出设备内存

# 4.4 绑定设备内存（输入→binding 0，输出→binding 1，与网络定义一致）
bindings = [int(d_input), int(d_output)]

# -------------------------- 5. 执行推理 --------------------------
# 5.1 将输入数据从主机拷贝到设备
cuda.memcpy_htod(d_input, input_data)  # htod: Host → Device

# 5.2 执行推理（记录时间）
start = time.time()
# 注意：execute_v2 要求 bindings 是设备内存地址列表
context.execute_v2(bindings=bindings)
end = time.time()

# 5.3 将输出数据从设备拷贝到主机（顺序：主机内存→设备内存）
cuda.memcpy_dtoh(output_data, d_output)  # dtoh: Device → Host

# -------------------------- 6. 结果输出与资源释放 --------------------------
print(f"推理耗时: {(end - start) * 1000:.2f} ms")
print(f"输入形状: {input_data.shape}, 输出形状: {output_data.shape}")
print(f"输出前5个值: {output_data[0][:5]}")

# 释放设备内存（可选，pycuda 会自动回收，但显式释放更规范）
d_input.free()
d_output.free()